import casadi as ca
import numpy as np

from src.model import Model
from .local_planner import LocalPlanner


class TrackingLocalPlanner(LocalPlanner):

    def __init__(self, model: Model, timestep, N=20):
        super().__init__(model, timestep, N)

        self._opt_u_ref = self._opti.parameter(N, 2)
        self._opt_x_ref = self._opti.parameter(N + 1, 3)

    def init(self):
        self._opti.subject_to(self._opt_states[0, :] == self._opt_x_ref[0, :])

        for i in range(self._N):
            state_error = self._opt_states[i, :] - self._opt_x_ref[i + 1, :]
            control_error = self._opt_controls[i, :] - self._opt_u_ref[i, :]
            self._cost += ca.mtimes([state_error, self._Q, state_error.T]) + ca.mtimes(
                [control_error, self._R, control_error.T])
        self._cost += self._slack_obs * ca.sumsqr(ca.vertcat(*self._slacks))
        self._opti.minimize(self._cost)

    def compute(self, next_trajectories, next_controls, obstacles):
        self._opti.set_value(self._opt_obstacles, obstacles)
        self._opti.set_value(self._opt_u_ref, next_controls)
        self._opti.set_value(self._opt_x_ref, next_trajectories)

        self._opti.set_initial(self._opt_controls, self._u0)
        self._opti.set_initial(self._opt_states, self._next_states)

        sol = self._opti.solve()

        u_res = sol.value(self._opt_controls)
        next_states_pred = sol.value(self._opt_states)

        self._u0 = np.concatenate((u_res[1:], u_res[-1:]))
        self._next_states = np.concatenate((next_states_pred[1:], next_states_pred[-1:]))

        return u_res[0]
